import numpy as np
import math
from machine_learning import train_model

#交易者类的基类
class Trader:
    def __init__(self, initial_money=0, tax_ratio=0.001, max_stock_num=10, stocks_info={}):
        self.cash = initial_money#手里剩下的钱
        self.stock_shares = {}#手里持有的股票份额
        for stock in stocks_info:
            self.stock_shares[stock] = 0
        self.net_worth = initial_money#净值
        self.tax_ratio = tax_ratio#手续费
        self.max_stock_num = max_stock_num#最多持有股票数
        self.stocks_info = stocks_info#各股票的历史信息

    #计算一下当前手中资产的总净值
    def get_worth(self, date):
        res = self.cash
        for stock_name in self.stock_shares:
            if not np.isnan(self.stocks_info[stock_name][date]['close']):
                res += self.stock_shares[stock_name] * self.stocks_info[stock_name][date]['close']
        return res

    #买入股票
    def buy(self, date, stock_name, money):
        if np.isnan(self.stocks_info[stock_name][date]['close']):
            raise ValueError(f'{date} {stock_name} \'s close value is nan')
        share = money / self.stocks_info[stock_name][date]['close']
        #买需要买整手
        share = share // 100 * 100
        cost = self.stocks_info[stock_name][date]['close'] * share
        #钱不够买就报错
        if cost > self.cash:
            raise ValueError(f'{date} {stock_name} not enough money')
        self.stock_shares[stock_name] += share
        self.cash -= cost

    #卖出股票
    def sell(self, date, stock_name, share):
        if np.isnan(self.stocks_info[stock_name][date]['close']):
            raise ValueError(f'{date} {stock_name} \'s close value is nan')
        #不够卖的就全卖光
        if share > self.stock_shares[stock_name]:
            money = self.stock_shares[stock_name] * self.stocks_info[stock_name][date]['close']
            self.stock_shares[stock_name] = 0
        else:
            money = share * self.stocks_info[stock_name][date]['close']
            self.stock_shares[stock_name] -= share
        # print(f'money:{res},date:{date}')
        self.cash += money * (1 - self.tax_ratio)

    #策略，在子类中实现
    def strategy(self, date):
        self.net_worth = self.get_worth(date)
        return
